import cv2
import numpy as np
import config as C
from fractal import multifractal_dimension as evaluate
import argparse
from utils import Image


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input_image')
    args = parser.parse_args()

    return args


def get_thresold():
    return cv2.getTrackbarPos('threshold', 'image') / 100.


def make_3channel(img):
    assert(len(img.shape) == 2)
    return np.stack([img, img, img], axis=2)


def get_output():
    threshold = get_thresold()
    bw_img = make_3channel(img.getBWImg())
    set_img = img.getSetImg()
    overlay = bw_img.copy()

    for x in range(0, width - C.BLOCK_WIDTH + 1, C.WIDTH_STEP):
        for y in range(0, height - C.BLOCK_HEIGHT + 1, C.HEIGHT_STEP):
            block = set_img[y:y + C.BLOCK_HEIGHT, x:x + C.BLOCK_WIDTH]
            evaluate_value = evaluate(block)
            if evaluate_value > threshold:
                overlay[y:y + C.BLOCK_HEIGHT, x:x + C.BLOCK_WIDTH, 2] = 255
                overlay[y:y + C.BLOCK_HEIGHT, x:x + C.BLOCK_WIDTH, 1] = 0
                overlay[y:y + C.BLOCK_HEIGHT, x:x + C.BLOCK_WIDTH, 0] = 0

    output = cv2.addWeighted(overlay, C.SHOW_ALPHA, bw_img, 1 - C.SHOW_ALPHA, 0)
    print('refreshed!')
    return output


def redraw(x):
    output = get_output()
    cv2.imshow('image', output)


if __name__ == '__main__':
    args = parse_args()
    img = Image(args.input_image, bw_threshold=C.GRAY2BW_THRESHOLD)
    width = img.data.shape[1]
    height = img.data.shape[0]
    cv2.namedWindow('image', cv2.WINDOW_NORMAL)
    cv2.createTrackbar('threshold', 'image', 0, 1, redraw)
    cv2.setTrackbarMin('threshold', 'image', 150)
    cv2.setTrackbarMax('threshold', 'image', 200)
    cv2.setTrackbarPos('threshold', 'image', 116)

    output = get_output()
    cv2.imshow('image', output)
    while(1):
        k = cv2.waitKey(0) & 0xFF
        if k == 27:
            break
    cv2.destroyAllWindows()
